#include <rclcpp/rclcpp.hpp>
#include <std_msgs/msg/string.hpp>
#include <opentelemetry/trace/provider.h>
#include <opentelemetry/exporters/ostream/span_exporter.h>
#include <opentelemetry/sdk/trace/simple_processor.h>
#include <opentelemetry/sdk/trace/tracer_provider.h>
#include <opentelemetry/sdk/trace/samplers/always_on.h>
#include <opentelemetry/sdk/resource/resource.h>
#include <opentelemetry/nostd/shared_ptr.h>
#include <opentelemetry/nostd/unique_ptr.h>

#include <chrono>
#include <thread>
#include <string>
#include <sstream>

namespace trace_api = opentelemetry::trace;
namespace trace_sdk = opentelemetry::sdk::trace;
namespace resource_sdk = opentelemetry::sdk::resource;
namespace nostd = opentelemetry::nostd;

// Initialize OpenTelemetry tracing
void InitTracer() {
    auto exporter = std::unique_ptr<trace_sdk::SpanExporter>(
        new opentelemetry::exporter::trace::OStreamSpanExporter);
    
    auto processor = std::unique_ptr<trace_sdk::SpanProcessor>(
        new trace_sdk::SimpleSpanProcessor(std::move(exporter)));
    
    auto resource = resource_sdk::Resource::Create({
        {"service.name", "ros2-publisher-node-with-trace"},
        {"service.version", "1.0.0"},
        {"service.instance.id", "publisher-with-trace-001"}
    });
    
    auto sampler = std::unique_ptr<trace_sdk::Sampler>(
        new trace_sdk::AlwaysOnSampler);
    
    auto provider = nostd::shared_ptr<trace_api::TracerProvider>(
        new trace_sdk::TracerProvider(std::move(processor), resource, std::move(sampler)));
    
    trace_api::Provider::SetTracerProvider(provider);
}

class PublisherNodeWithTrace : public rclcpp::Node {
public:
    PublisherNodeWithTrace() : Node("publisher_node_with_trace") {
        // Initialize OpenTelemetry
        InitTracer();
        tracer_ = trace_api::Provider::GetTracerProvider()->GetTracer("publisher-node-with-trace");
        
        // Create publisher
        publisher_ = this->create_publisher<std_msgs::msg::String>("data_topic_with_trace", 10);
        
        // Create timer to publish messages
        timer_ = this->create_wall_timer(
            std::chrono::milliseconds(3000),
            std::bind(&PublisherNodeWithTrace::publishMessage, this));
        
        RCLCPP_INFO(this->get_logger(), "Publisher Node with Trace started");
        RCLCPP_INFO(this->get_logger(), "This node will pass trace_id and span_id to next nodes");
    }

private:
    void publishMessage() {
        // Create root span for message publishing
        auto span = tracer_->StartSpan("publish_data_message_with_trace", {
            {"node.name", "publisher_node_with_trace"},
            {"topic.name", "data_topic_with_trace"},
            {"message.type", "std_msgs/String"},
            {"operation.type", "publish"},
            {"message.count", static_cast<int64_t>(message_count_)}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            // Get trace context information (using simple string representation)
            auto context = span->GetContext();
            std::string trace_id = "trace_" + std::to_string(message_count_);
            std::string span_id = "span_pub_" + std::to_string(message_count_);
            
            // Create message with embedded trace information
            auto message = std_msgs::msg::String();
            std::stringstream ss;
            ss << "trace_id:" << trace_id 
               << "|span_id:" << span_id 
               << "|data:Data from Publisher with Trace - " << message_count_;
            message.data = ss.str();
            
            // Add span attributes
            span->SetAttribute("message.content", message.data);
            span->SetAttribute("trace.id", trace_id);
            span->SetAttribute("span.id", span_id);
            span->AddEvent("message_created_with_trace");
            
            // Simulate some processing time
            std::this_thread::sleep_for(std::chrono::milliseconds(30));
            
            // Publish message
            publisher_->publish(message);
            span->AddEvent("message_published_with_trace");
            
            // Set success status
            span->SetStatus(trace_api::StatusCode::kOk, "Message published with trace successfully");
            
            RCLCPP_INFO(this->get_logger(), "Published with trace: %s", message.data.c_str());
            RCLCPP_INFO(this->get_logger(), "Trace ID: %s, Span ID: %s", trace_id.c_str(), span_id.c_str());
            
            message_count_++;
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            RCLCPP_ERROR(this->get_logger(), "Error publishing message: %s", e.what());
        }
        
        span->End();
    }

    rclcpp::Publisher<std_msgs::msg::String>::SharedPtr publisher_;
    rclcpp::TimerBase::SharedPtr timer_;
    nostd::shared_ptr<trace_api::Tracer> tracer_;
    size_t message_count_ = 0;
};

int main(int argc, char* argv[]) {
    rclcpp::init(argc, argv);
    
    auto node = std::make_shared<PublisherNodeWithTrace>();
    
    RCLCPP_INFO(node->get_logger(), "Starting Publisher Node with Cross-Node Tracing...");
    RCLCPP_INFO(node->get_logger(), "This demo shows how to pass trace_id and span_id across nodes");
    
    rclcpp::spin(node);
    rclcpp::shutdown();
    
    return 0;
}
